#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import rclpy
from rclpy.node import Node
from rclpy.duration import Duration
import numpy as np
import tf2_ros
from geometry_msgs.msg import WrenchStamped
from std_srvs.srv import Trigger

def quat_to_rotmat(x, y, z, w):
    n = np.sqrt(x*x + y*y + z*z + w*w)
    x, y, z, w = x/n, y/n, z/n, w/n
    xx, yy, zz = x*x, y*y, z*z
    xy, xz, yz = x*y, x*z, y*z
    wx, wy, wz = w*x, w*y, w*z
    return np.array([
        [1 - 2*(yy+zz), 2*(xy-wz),     2*(xz+wy)],
        [2*(xy+wz),     1 - 2*(xx+zz), 2*(yz-wx)],
        [2*(xz-wy),     2*(yz+wx),     1 - 2*(xx+yy)]
    ], dtype=float)

class GravityCalibNode(Node):
    def __init__(self):
        super().__init__('gravity_calib_node')

        # 参数
        self.declare_parameter('base_frame', 'base_link')
        self.declare_parameter('ee_frame', 'end')
        self.declare_parameter('wrench_topic', '/force_data')
        self.declare_parameter('capture_timeout_s', 2.0)

        self.base_frame   = self.get_parameter('base_frame').value
        self.ee_frame     = self.get_parameter('ee_frame').value
        self.wrench_topic = self.get_parameter('wrench_topic').value
        self.cap_timeout  = self.get_parameter('capture_timeout_s').value

        # 末端->力传感器固定旋转
        self.M = np.array([[0, 0, 1],
                           [0,-1, 0],
                           [1, 0, 0]], dtype=float)

        # TF
        self.tf_buffer = tf2_ros.Buffer()
        self.tf_listener = tf2_ros.TransformListener(self.tf_buffer, self)

        # 订阅力传感器
        self.last_force = None
        self.sub_wrench = self.create_subscription(WrenchStamped, self.wrench_topic, self.cb_wrench, 10)

        # 发布补偿后的力
        self.pub_comp = self.create_publisher(WrenchStamped, '/wrench_comp', 10)

        # 样本数据
        self.samples_RSB = []
        self.samples_F   = []

        self.g_B = None
        self.F0  = None

        # 服务
        self.srv_capture = self.create_service(Trigger, 'capture_sample', self.handle_capture)
        self.srv_solve   = self.create_service(Trigger, 'solve_calibration', self.handle_solve)

        # 定时器：补偿并发布
        self.timer = self.create_timer(0.01, self.online_compensation)

        self.get_logger().info(f'GravityCalibNode up. base={self.base_frame}, ee={self.ee_frame}, wrench={self.wrench_topic}')

    def cb_wrench(self, msg: WrenchStamped):
        self.last_force = np.array([msg.wrench.force.x,
                                    msg.wrench.force.y,
                                    msg.wrench.force.z], dtype=float)

    def lookup_R_SB(self) -> np.ndarray:
        if not self.tf_buffer.can_transform(self.base_frame, self.ee_frame, rclpy.time.Time(),
                                            timeout=Duration(seconds=self.cap_timeout)):
            raise RuntimeError(f'等待 TF 超时: {self.base_frame}->{self.ee_frame}')
        trans = self.tf_buffer.lookup_transform(self.base_frame, self.ee_frame, rclpy.time.Time(),
                                                timeout=Duration(seconds=0.5))
        q = trans.transform.rotation
        R_BE = quat_to_rotmat(q.x, q.y, q.z, q.w)
        R_BS = R_BE @ self.M
        R_SB = R_BS.T
        return R_SB

    def handle_capture(self, request, response):
        try:
            if self.last_force is None:
                raise RuntimeError('暂未收到力传感器数据')
            R_SB = self.lookup_R_SB()
            self.samples_RSB.append(R_SB)
            self.samples_F.append(self.last_force.copy())

            n = len(self.samples_F)
            response.success = True
            response.message = f'采样成功，第 {n} 组。'
            self.get_logger().info(response.message)
        except Exception as e:
            response.success = False
            response.message = f'采样失败: {e}'
            self.get_logger().warn(response.message)
        return response

    def handle_solve(self, request, response):
        try:
            N = len(self.samples_F)
            if N < 3:
                raise RuntimeError(f'样本不足，需要≥3，当前 {N}')

            A_rows, b_rows = [], []
            for i in range(N):
                R = self.samples_RSB[i]
                F = self.samples_F[i]
                A_rows.append(np.hstack([R, np.eye(3)]))
                b_rows.append(F)
            A = np.vstack(A_rows)
            b = np.hstack(b_rows)

            x, *_ = np.linalg.lstsq(A, b, rcond=None)
            self.g_B = x[0:3]
            self.F0  = x[3:6]

            g_norm = np.linalg.norm(self.g_B)
            self.get_logger().info(f'求解完成: g_B={self.g_B} (|g_B|={g_norm:.4f}), F0={self.F0}')
            response.success = True
            response.message = '标定完成，进入在线补偿模式。'
        except Exception as e:
            response.success = False
            response.message = f'求解失败: {e}'
            self.get_logger().error(response.message)
        return response

    def online_compensation(self):
        if self.g_B is None or self.F0 is None or self.last_force is None:
            return
        try:
            R_SB = self.lookup_R_SB()
            F_pred = R_SB @ self.g_B + self.F0
            F_comp = self.last_force - F_pred

            msg = WrenchStamped()
            msg.header.stamp = self.get_clock().now().to_msg()
            msg.header.frame_id = 'sensor'  # 你可改成真实 frame
            msg.wrench.force.x = float(F_comp[0])
            msg.wrench.force.y = float(F_comp[1])
            msg.wrench.force.z = float(F_comp[2])
            self.pub_comp.publish(msg)
        except Exception as e:
            self.get_logger().warn(f'在线补偿失败: {e}')

def main():
    rclpy.init()
    node = GravityCalibNode()
    rclpy.spin(node)
    node.destroy_node()
    rclpy.shutdown()

if __name__ == '__main__':
    main()
